import torch
from torch import nn
from d2l import torch as d2l

# 实现互相关运算
def corr2d(X, K):
    h, w = K.shape
    Y = torch.zeros(X.shape[0]-h+1, X.shape[1]-w+1)
    for i in range(Y.shape[0]):
        for j in range(Y.shape[1]):
            Y[i, j] = (X[i:i+h, j:j+w]*K).sum()
    return Y
# 验证计算是否正确
X = torch.tensor([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]])
K = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
corr2d(X, K)

# 基于上面定义的corr2d函数实现二维卷积层。在__init__构造函数中，将weight和bias声明为两个模型参数。 前向传播函数调用corr2d函数并添加偏置。
class Conv2D(nn.Module):
    def __init__(self, kernel_size):
        super(Conv2D, self).__init__()
        self.weight = nn.Parameter(torch.randn(kernel_size))  # 记住参数定义不要忘记了nn的Parameter
        self.bias = nn.Parameter(torch.zerps(1))

    def forward(self, X):
        return corr2d(X, self.weight) + self.bias

# 如下是卷积层的一个简单应用：通过找到像素变化的位置，来检测图像中不同颜色的边缘。
# 首先，我们构造 一个6 × 8像素的黑白图像。中间四列为黑色（0），其余像素为白色（1）
X = torch.ones(6, 8)
X[:, 2:6] = 0
X

# 接下来，我们构造一个高度为1、宽度为2的卷积核K。当进行互相关运算时，如果水平相邻的两元素相同，则 输出为零，否则输出为非零。
K = torch.tensor([[1.0, -1.0]])
Y = corr2d(X, K)
Y
# 现在我们将输入的二维图像转置，再进行如上的互相关运算
corr2d(X.t(), K)  # 不出所料，这个卷积核K只可以检测垂直边缘，无法检测水平边缘。

# 学习卷积核
# 为了简单起见，我们在此使用内置的二维卷积层，并忽略偏置。
conv2d = nn.Conv2d(1, 1, kernel_size=(1, 2), bias=False)  # 这里的1，1表示输入和输出通道

# 这个二维卷积层使用四维输入和输出格式（批量大小、通道、高度、宽度），其中批量大小和通道数都为1
X = X.reshape((1, 1, 6, 8))
Y = Y.reshape((1, 1, 6, 7))
lr = 3e-2  # 学习率

# 训练
for i in range(20):
    Y_hat = conv2d(X)  # 可知卷积层不在意输入的高度和宽度
    l = (Y_hat - Y) ** 2
    conv2d.zero_grad()
    l.sum().backward()
    # 迭代卷积核
    conv2d.weight.data[:] -= lr * conv2d.weight.grad
    if (i + 1) % 2 == 0:
        print(f'epoch {i + 1}, loss {l.sum():.3f}')

# 查看一下训练后的权重
conv2d.weight.data.reshape(1, 2)